package s9_树.tree4_赫夫曼树;

import java.util.*;

/**
 * <code>{@link HuffmanTree}</code>
 * <p>
 * description: Tree
 * <p>
 *
 * @author 西瓜瓜 on 2022/2/23 20:05
 *
 * 赫夫曼编码
 *
 */
public class HuffmanCode {

    public static void main(String[] args) {
        String content = "i like like like java do you like a java";
        byte[] contentBytes = content.getBytes();
        System.out.println(Arrays.toString(contentBytes));

        HuffmanCode huffmanCode = new HuffmanCode();
        List<CNode> nodes = huffmanCode.getNodes(contentBytes);

        CNode huffmanTree = huffmanCode.createHuffmanTree(nodes);

        Map<Byte, String> huffmanCodeMap = huffmanCode.huffmanCode(huffmanTree);

        byte[] huffmanCodeData = huffmanCode.huffmanCodeData(contentBytes, huffmanCodeMap);
//        System.out.println(Arrays.toString(huffmanCodeData));

        byte[] bytes = huffmanCode.decodeHuffmanCodeData(huffmanCodeData, huffmanCodeMap);
        System.out.println(Arrays.toString(bytes));
        System.out.println(new String(bytes));



    }


    /**
     * 返回压缩后的byte数组
     * @param contentBytes 原数组
     * @param huffmanCodeMap 赫夫曼转码表
     * @return
     */
    private byte[] huffmanCodeData(byte[] contentBytes, Map<Byte, String> huffmanCodeMap) {
        StringBuilder builder = new StringBuilder();
        for (byte contentByte : contentBytes) {
            builder.append(huffmanCodeMap.get(contentByte));
        }
        // 统计返回byte[] 长度
        int len = (builder.length() % 8 == 0) ? (builder.length() / 8) : (builder.length() / 8 + 1);
        // 创建存储压缩后的byte数组
        byte[] huffmanCodeBytes = new byte[len];
        // 每8位对应一个byte，所以步长+8
        int index = 0;
        for(int i=0; i<builder.length(); i += 8) {
            int end = ((i+8) > builder.length()) ? builder.length() : (i + 8);
            String strByte = builder.substring(i, end);
//            System.out.println(strByte);
            // 将strByte转成byte
            huffmanCodeBytes[index++] = (byte) Integer.parseInt(strByte, 2);
        }
        return huffmanCodeBytes;
    }


    /**
     * 将压缩后的byte[]还原转换成原始的byte[]
     * @param huffmanCodeBytes
     * @param huffmanCodeMap
     * @return
     */
    private byte[] decodeHuffmanCodeData(byte[] huffmanCodeBytes, Map<Byte, String> huffmanCodeMap) {
        StringBuilder builder = new StringBuilder();
        for (int i = 0; i < huffmanCodeBytes.length; i++) {
            String bitString = byteToBitString(i < huffmanCodeBytes.length - 1, huffmanCodeBytes[i]);
            builder.append(bitString);
        }

        Map<String, Byte> sourceMap = new HashMap<>(huffmanCodeMap.size());
        huffmanCodeMap.entrySet().forEach(entry -> sourceMap.put(entry.getValue(), entry.getKey()));

        List<Byte> sourceByteList = new ArrayList<>();
        int start = 0;
        int end = 0;
        while(end < builder.length()) {
            end++;
            String substring = builder.substring(start, end);
            if(sourceMap.containsKey(substring)) {
                sourceByteList.add(sourceMap.get(substring));
                start = end;
            }
        }
        if(start != end) {
            String substring = builder.substring(start, end);
            if(sourceMap.containsKey(substring)) {
                sourceByteList.add(sourceMap.get(substring));
            }
        }

        byte[] bytes = new byte[sourceByteList.size()];
        for (int i = 0; i < sourceByteList.size(); i++) {
            bytes[i] = sourceByteList.get(i);
        }
        return bytes;
    }

    private String byteToBitString(boolean flag, int huffmanCodeByte) {
        int temp = huffmanCodeByte;
        if(flag) {
            temp |= 256;
        }
        String str = Integer.toBinaryString(temp);
        if(flag) {
            return str.substring(str.length() - 8);
        } else {
            return str;
        }
    }

    /**
     * 创建赫夫曼编码表
     *
     * @param huffmanTree
     * @return
     */
    private Map<Byte, String> huffmanCode(CNode huffmanTree) {
        Map<Byte, String> huffmanCodeMap = new HashMap<>();
        // 前序遍历
        huffmanTree.preOrder(huffmanCodeMap, "");
        return huffmanCodeMap;
    }

    private List<CNode> getNodes(byte[] bytes) {
        List<CNode> nodes = new ArrayList<>();

        // 统计每一个字节出现的次数
        Map<Byte, Integer> byteCount = new HashMap<>();
        for (byte b : bytes) {
            Integer count = byteCount.getOrDefault(b, 0);
            count += 1;
            byteCount.put(b, count);
        }

        // 把每一个Entry转换成CNode对象
        for (Map.Entry<Byte, Integer> entry : byteCount.entrySet()) {
            nodes.add(new CNode(entry.getKey(), entry.getValue()));
        }

        return nodes;
    }

    /**
     * 创建赫夫曼树
     *
     * @param nodes
     * @return
     */
    private CNode createHuffmanTree(List<CNode> nodes) {
        while(nodes.size() > 1) {
            Collections.sort(nodes);

            CNode leftNode = nodes.get(0);
            CNode rightNode = nodes.get(1);
            CNode parentNode = new CNode(leftNode.weight + rightNode.weight);
            parentNode.leftNode = leftNode;
            parentNode.rightNode = rightNode;

            nodes.remove(leftNode);
            nodes.remove(rightNode);

            nodes.add(parentNode);
        }

        return nodes.get(0);
    }
}

class CNode implements Comparable<CNode>{
    /** 存放数据的本身 */
    Byte data;
    /** 节点的权值,表示字符出现的次数 */
    int weight;
    /** 左子节点 */
    CNode leftNode;
    /** 右子节点 */
    CNode rightNode;

    /**
     * 前序遍历
     */
    public void preOrder() {
        System.out.println(this);
        if(this.leftNode != null) {
            this.leftNode.preOrder();
        }
        if(this.rightNode != null) {
            this.rightNode.preOrder();
        }
    }

    public void preOrder(Map<Byte, String> huffmanCodeMap, String code) {
        if(this.data != null) {
            huffmanCodeMap.put(this.data, code);
        }

        if(this.leftNode != null) {
            this.leftNode.preOrder(huffmanCodeMap, code+"0");
        }
        if(this.rightNode != null) {
            this.rightNode.preOrder(huffmanCodeMap, code+"1");
        }

    }

    public CNode(int weight) {
        this.weight = weight;
    }

    public CNode(Byte data, int weight) {
        this.data = data;
        this.weight = weight;
    }

    @Override
    public String toString() {
        return "CNode{" +
                "data=" + data +
                ", weight=" + weight +
                '}';
    }

    @Override
    public int compareTo(CNode o) {
        return this.weight - o.weight;
    }


}
